#!/usr/bin/env python3
"""
User Plane Interface Analysis Script
Automatically analyzes all user plane PCAP files and generates mermaid charts
"""

import json
import os
import sys
from pathlib import Path
import argparse
from datetime import datetime

class UserPlaneAnalyzer:
    def __init__(self, captures_dir="../docs/5gc-interfaces/captures/user-plane"):
        self.captures_dir = Path(captures_dir)
        self.output_dir = Path("./analysis/user-plane-interfaces")
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
    def analyze_pcap_file(self, pcap_file):
        """Analyze a single PCAP file and return analysis results"""
        try:
            # Import the PCAP to JSON converter
            import subprocess
            import tempfile
            
            # Use the pcap-to-json.py script directly
            json_file = pcap_file.with_suffix('.json')
            
            # Run the conversion
            result = subprocess.run([
                'python3', 'pcap-to-json.py', 
                str(pcap_file), str(json_file)
            ], capture_output=True, text=True, cwd='.')
            
            if result.returncode != 0:
                print(f"Conversion failed: {result.stderr}")
                return None
            
            # Load the JSON data
            with open(json_file, 'r') as f:
                data = json.load(f)
            
            return data
            
        except Exception as e:
            print(f"Error analyzing {pcap_file}: {e}")
            return None
    
    def generate_interface_summary(self, data, interface_name):
        """Generate summary for an interface"""
        if not data:
            return None
            
        summary = {
            'interface': interface_name,
            'total_packets': data['metadata']['total_packets'],
            'protocols': data['statistics']['protocols'],
            'source_ips': data['statistics']['source_ips'],
            'dest_ips': data['statistics']['dest_ips'],
            'ports': data['statistics']['ports']
        }
        
        return summary
    
    def generate_mermaid_chart(self, interface_name, summary):
        """Generate mermaid chart for an interface"""
        if not summary:
            return ""
            
        # Define interface-specific configurations
        interface_configs = {
            'n1': {
                'description': 'N1 Interface (UE ↔ gNodeB)',
                'protocol': 'NAS over Radio',
                'participants': ['UE', 'gNodeB'],
                'flow': 'UE ↔ gNodeB (Radio)'
            },
            'n3': {
                'description': 'N3 Interface (gNodeB ↔ UPF)',
                'protocol': 'GTP-U over UDP',
                'participants': ['gNodeB', 'UPF'],
                'flow': 'gNodeB ↔ UPF (GTP-U)'
            },
            'n6': {
                'description': 'N6 Interface (UPF ↔ Internet)',
                'protocol': 'IP Routing',
                'participants': ['UPF', 'Internet'],
                'flow': 'UPF ↔ Internet (NAT)'
            }
        }
        
        config = interface_configs.get(interface_name.lower(), {
            'description': f'{interface_name.upper()} Interface',
            'protocol': 'Unknown',
            'participants': ['Source', 'Destination'],
            'flow': 'Data Flow'
        })
        
        # Generate mermaid sequence diagram
        mermaid = f"""## {config['description']}

### Interface Summary:
- **Total Packets**: {summary['total_packets']}
- **Protocols**: {', '.join(summary['protocols'].keys())}
- **Source IPs**: {', '.join(list(summary['source_ips'].keys())[:3])}
- **Destination IPs**: {', '.join(list(summary['dest_ips'].keys())[:3])}

### Protocol Stack Flow:

```mermaid
sequenceDiagram
    participant UE as User Equipment
    participant GNB as gNodeB<br/>(192.168.1.221)
    participant UPF as UPF<br/>(192.168.1.219)
    participant INT as Internet
    
    Note over UE,INT: {config['description']}
    
"""
        
        # Add specific flow based on interface
        if interface_name.lower() == 'n1':
            mermaid += """    UE->>GNB: NAS Messages (Radio)
    Note right of GNB: RRC/PDCP/RLC/MAC/PHY Processing
    GNB->>UE: NAS Response (Radio)
    Note right of UE: Application Processing
"""
        elif interface_name.lower() == 'n3':
            mermaid += """    GNB->>UPF: GTP-U Encapsulated Data<br/>Src: 192.168.1.221:2152<br/>Dst: 192.168.1.219:2152
    Note right of UPF: GTP-U Decapsulation
    UPF->>GNB: GTP-U Encapsulated Data<br/>Src: 192.168.1.219:2152<br/>Dst: 192.168.1.221:2152
    Note right of GNB: GTP-U Processing
"""
        elif interface_name.lower() == 'n6':
            mermaid += """    UPF->>INT: User Data (NAT)<br/>Src: 192.168.1.219<br/>Dst: Internet IP
    Note right of INT: Internet Processing
    INT->>UPF: Response Data (NAT)<br/>Src: Internet IP<br/>Dst: 192.168.1.219
    Note right of UPF: NAT Translation
"""
        
        mermaid += """```

### Traffic Analysis:
"""
        
        # Add protocol analysis
        for protocol, count in summary['protocols'].items():
            mermaid += f"- **{protocol}**: {count} packets\n"
        
        # Add IP analysis
        mermaid += f"\n### IP Address Analysis:\n"
        for ip, count in list(summary['source_ips'].items())[:3]:
            mermaid += f"- **{ip}**: {count} packets (source)\n"
        for ip, count in list(summary['dest_ips'].items())[:3]:
            mermaid += f"- **{ip}**: {count} packets (destination)\n"
        
        return mermaid
    
    def analyze_all_interfaces(self):
        """Analyze all user plane interfaces"""
        print("=== User Plane Interface Analysis ===")
        
        # Find all PCAP files
        pcap_files = list(self.captures_dir.glob("*.pcap"))
        
        if not pcap_files:
            print(f"No PCAP files found in {self.captures_dir}")
            return
        
        print(f"Found {len(pcap_files)} PCAP files")
        
        # Analyze each file
        results = {}
        for pcap_file in pcap_files:
            interface_name = pcap_file.stem.replace('_', ' ').title()
            print(f"\nAnalyzing {interface_name}...")
            
            data = self.analyze_pcap_file(pcap_file)
            if data:
                summary = self.generate_interface_summary(data, interface_name)
                results[interface_name] = summary
                print(f"✓ {interface_name}: {summary['total_packets']} packets")
            else:
                print(f"✗ {interface_name}: Analysis failed")
        
        # Generate comprehensive report
        self.generate_comprehensive_report(results)
        
        return results
    
    def generate_comprehensive_report(self, results):
        """Generate comprehensive analysis report"""
        report_file = self.output_dir / "user-plane-analysis-report.md"
        
        with open(report_file, 'w') as f:
            f.write("# User Plane Interface Analysis Report\n\n")
            f.write(f"Generated: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
            
            f.write("## Overview\n\n")
            f.write("This report provides detailed analysis of all user plane interfaces in the 5GC testbed.\n\n")
            
            # Summary table
            f.write("## Interface Summary\n\n")
            f.write("| Interface | Total Packets | Protocols | Status |\n")
            f.write("|-----------|---------------|-----------|--------|\n")
            
            for interface_name, summary in results.items():
                if summary:
                    protocols = ', '.join(summary['protocols'].keys())
                    f.write(f"| {interface_name} | {summary['total_packets']} | {protocols} | ✅ Active |\n")
                else:
                    f.write(f"| {interface_name} | 0 | None | ❌ Failed |\n")
            
            f.write("\n## Detailed Analysis\n\n")
            
            # Generate mermaid charts for each interface
            for interface_name, summary in results.items():
                if summary:
                    mermaid_chart = self.generate_mermaid_chart(interface_name, summary)
                    f.write(mermaid_chart)
                    f.write("\n---\n\n")
        
        print(f"\nComprehensive report generated: {report_file}")
    
    def generate_combined_mermaid(self, results):
        """Generate combined mermaid chart for all interfaces"""
        mermaid_file = self.output_dir / "user-plane-combined-flow.md"
        
        with open(mermaid_file, 'w') as f:
            f.write("# Combined User Plane Flow\n\n")
            f.write("## Complete 5G User Plane Architecture\n\n")
            
            f.write("```mermaid\n")
            f.write("graph TB\n")
            f.write("    subgraph 'User Equipment (UE)'\n")
            f.write("        UE[UE Application]\n")
            f.write("        UE_NAS[UE NAS Layer]\n")
            f.write("    end\n\n")
            
            f.write("    subgraph 'gNodeB (192.168.1.221)'\n")
            f.write("        GNB_RRC[gNodeB RRC]\n")
            f.write("        GNB_PDCP[gNodeB PDCP]\n")
            f.write("        GNB_GTPU[gNodeB GTP-U]\n")
            f.write("    end\n\n")
            
            f.write("    subgraph 'UPF (192.168.1.219)'\n")
            f.write("        UPF_GTPU[UPF GTP-U]\n")
            f.write("        UPF_NAT[UPF NAT]\n")
            f.write("    end\n\n")
            
            f.write("    subgraph 'Internet'\n")
            f.write("        INT[Internet Services]\n")
            f.write("    end\n\n")
            
            # N1 Interface
            f.write("    %% N1 Interface (UE ↔ gNodeB)\n")
            f.write("    UE --> UE_NAS\n")
            f.write("    UE_NAS -.->|N1: Radio| GNB_RRC\n")
            f.write("    GNB_RRC --> GNB_PDCP\n\n")
            
            # N3 Interface
            f.write("    %% N3 Interface (gNodeB ↔ UPF)\n")
            f.write("    GNB_PDCP --> GNB_GTPU\n")
            f.write("    GNB_GTPU -.->|N3: GTP-U| UPF_GTPU\n\n")
            
            # N6 Interface
            f.write("    %% N6 Interface (UPF ↔ Internet)\n")
            f.write("    UPF_GTPU --> UPF_NAT\n")
            f.write("    UPF_NAT -.->|N6: NAT| INT\n\n")
            
            # Downlink flow
            f.write("    %% Downlink Flow\n")
            f.write("    INT -.->|N6: Response| UPF_NAT\n")
            f.write("    UPF_NAT --> UPF_GTPU\n")
            f.write("    UPF_GTPU -.->|N3: GTP-U| GNB_GTPU\n")
            f.write("    GNB_GTPU --> GNB_PDCP\n")
            f.write("    GNB_PDCP --> GNB_RRC\n")
            f.write("    GNB_RRC -.->|N1: Radio| UE_NAS\n")
            f.write("    UE_NAS --> UE\n\n")
            
            # Styling
            f.write("    %% Styling\n")
            f.write("    classDef gNodeB fill:#e1f5fe\n")
            f.write("    classDef UPF fill:#f3e5f5\n")
            f.write("    classDef UE fill:#e8f5e8\n")
            f.write("    classDef internet fill:#fff3e0\n\n")
            
            f.write("    class GNB_RRC,GNB_PDCP,GNB_GTPU gNodeB\n")
            f.write("    class UPF_GTPU,UPF_NAT UPF\n")
            f.write("    class UE,UE_NAS UE\n")
            f.write("    class INT internet\n")
            f.write("```\n")
        
        print(f"Combined mermaid chart generated: {mermaid_file}")

def main():
    parser = argparse.ArgumentParser(description='Analyze user plane interfaces')
    parser.add_argument('--captures-dir', default='../docs/5gc-interfaces/captures/user-plane',
                       help='Directory containing PCAP files')
    parser.add_argument('--output-dir', default='./analysis/user-plane-interfaces',
                       help='Output directory for analysis')
    
    args = parser.parse_args()
    
    analyzer = UserPlaneAnalyzer(args.captures_dir)
    results = analyzer.analyze_all_interfaces()
    
    if results:
        analyzer.generate_combined_mermaid(results)
        print("\n=== Analysis Complete ===")
        print("Generated files:")
        print(f"  - {analyzer.output_dir}/user-plane-analysis-report.md")
        print(f"  - {analyzer.output_dir}/user-plane-combined-flow.md")
    else:
        print("No results generated")

if __name__ == '__main__':
    main() 